# -*- coding: utf-8 -*-
# @Time    : 2025/4/6 21:45
# @Author  : sjh
# @Site    : 
# @File    : Stereo.py
# @Comment :
import cv2
import numpy as np
import json
import glob
import os
from stereo.utils.config import Config
# np.set_printoptions(suppress=True)

# 全局变量存储校准参数
_initialized = False
_image_size = Config.image_size
_k_left = None
_k_right = None
_dist_left = None
_dist_right = None
_R = None
_T = None
_Q = None
_map1x = None
_map1y = None
_map2x = None
_map2y = None
_fx = None
_fy = None
_cx = None
_cy = None
_baseline = None
_depth_cam_matrix = None

# NPZ文件路径，用于存储计算结果
_NPZ_CACHE_PATH = os.path.join('stereo', 'calib', 'stereo_rectify_cache.npz')

def scale_camera_matrix(K, orig_size, new_size):
    scale_x = new_size[0] / orig_size[0]
    scale_y = new_size[1] / orig_size[1]
    K_scaled = np.array([
        [K[0, 0] * scale_x, 0, K[0, 2] * scale_x],
        [0, K[1, 1] * scale_y, K[1, 2] * scale_y],
        [0, 0, 1]
    ], dtype=np.float64)
    return K_scaled

def initialize(resolution=(640, 480), force_recalculate=False):
    """初始化立体校正参数，如果有缓存文件则直接加载"""
    global _initialized, _image_size, _k_left, _k_right, _dist_left, _dist_right
    global _R, _T, _Q, _map1x, _map1y, _map2x, _map2y
    global _fx, _fy, _cx, _cy, _baseline, _depth_cam_matrix
    
    # 如果已初始化，除非强制重新计算，否则直接返回
    if _initialized and not force_recalculate:
        return
    
    # 首先尝试从缓存加载
    if not force_recalculate and os.path.exists(_NPZ_CACHE_PATH):
        try:
            print(f"从缓存加载立体校正参数: {_NPZ_CACHE_PATH}")
            cache = np.load(_NPZ_CACHE_PATH, allow_pickle=True)
            
            _k_left = cache['k_left']
            _k_right = cache['k_right']
            _dist_left = cache['dist_left']
            _dist_right = cache['dist_right']
            _R = cache['R']
            _T = cache['T']
            _Q = cache['Q']
            _map1x = cache['map1x']
            _map1y = cache['map1y']
            _map2x = cache['map2x']
            _map2y = cache['map2y']
            _fx = float(cache['fx'])
            _fy = float(cache['fy'])
            _cx = float(cache['cx'])
            _cy = float(cache['cy'])
            _baseline = float(cache['baseline'])
            _depth_cam_matrix = cache['depth_cam_matrix']
            
            # 如果分辨率不同，需要重新计算
            saved_resolution = tuple(cache['resolution'])
            if saved_resolution != resolution:
                print(f"分辨率不匹配，重新计算参数: 缓存={saved_resolution}, 请求={resolution}")
                force_recalculate = True
            else:
                _initialized = True
                print("立体校正参数加载成功")
                print(f"fx={_fx}, fy={_fy}, cx={_cx}, cy={_cy}, baseline={_baseline}")
                return
        except Exception as e:
            print(f"加载缓存失败: {e}")
            force_recalculate = True
    
    # 如果需要计算参数
    print("重新计算立体校正参数...")
    json_path = os.path.join('stereo', 'calib', 'cali_circle.json')
    k_left, dist_left, k_right, dist_right, R, T, Q = load_calibration_from_json(json_path)
    _k_left, _k_right = np.array(k_left), np.array(k_right)
    
    _k_left = scale_camera_matrix(_k_left, _image_size, (640, 480))
    _k_right = scale_camera_matrix(_k_right, _image_size, (640, 480))
    
    
    _dist_left, _dist_right = np.array(dist_left), np.array(dist_right)
    _R, _T, _Q = np.array(R), np.array(T), np.array(Q)
    
    # 计算映射
    _map1x, _map1y, _map2x, _map2y, _ = stereo_rectify_with_distortion(
        _k_left, _dist_left, _k_right, _dist_right, _R, _T
    )
    
    # 计算内参和基线
    calculate_fx_fy_cx_cy_baseline(_Q)
    
    # 根据分辨率调整内参
    ori_width, ori_height = _image_size  # 原始图像
    scale_x = resolution[0] / ori_width  # 宽度缩放比例
    scale_y = resolution[1] / ori_height  # 高度缩放比例
    reset_calib(scale_x, scale_y)  # 调整内参
    
    _initialized = True
    print("立体校正参数计算完成")
    print(f"fx={_fx}, fy={_fy}, cx={_cx}, cy={_cy}, baseline={_baseline}")
    
    # 保存缓存
    try:
        os.makedirs(os.path.dirname(_NPZ_CACHE_PATH), exist_ok=True)
        np.savez(_NPZ_CACHE_PATH,
                 k_left=_k_left,
                 k_right=_k_right,
                 dist_left=_dist_left,
                 dist_right=_dist_right,
                 R=_R,
                 T=_T,
                 Q=_Q,
                 map1x=_map1x,
                 map1y=_map1y,
                 map2x=_map2x,
                 map2y=_map2y,
                 fx=_fx,
                 fy=_fy,
                 cx=_cx,
                 cy=_cy,
                 baseline=_baseline,
                 depth_cam_matrix=_depth_cam_matrix,
                 resolution=resolution)
        print(f"立体校正参数已缓存至: {_NPZ_CACHE_PATH}")
    except Exception as e:
        print(f"保存缓存失败: {e}")

def stereo_rectify_with_distortion(K_left, dist_left, K_right, dist_right, R, T, newimage_size=None):
    """计算有畸变立体校正变换"""
    # 图像尺寸
    gdc_width, gdc_height = _image_size  # 原始图像尺寸
    out_width, out_height = Config.image_size  # 校正后图像尺寸

    # 调用stereoRectify
    flags = cv2.CALIB_ZERO_DISPARITY  # 强制校正后主点垂直对齐
    alpha = 0  # 自动裁剪黑边
    
    R1, R2, P1, P2, Q, roi1, roi2 = cv2.stereoRectify(
            cameraMatrix1=K_left,  # 左相机内参
            distCoeffs1=dist_left,  # 左相机畸变系数
            cameraMatrix2=K_right,  # 右相机内参
            distCoeffs2=dist_right,  # 右相机畸变系数
            imageSize=(gdc_width, gdc_height),  # 原始图像尺寸
            R=R,  # 旋转矩阵 右到左的旋转矩阵
            T=T,  # 平移矩阵 右到左的平移向量
            flags=flags,  # 标志
            alpha=alpha,  # 自动裁剪黑边
            newImageSize=(out_width, out_height)  # 校正后图像尺寸
        )
    # 2. 计算映射表（考虑畸变）
    map1x, map1y = cv2.initUndistortRectifyMap(K_left, dist_left, R1, P1, _image_size, cv2.CV_32FC1)
    map2x, map2y = cv2.initUndistortRectifyMap(K_right, dist_right, R2, P2, _image_size, cv2.CV_32FC1)
    
    return map1x, map1y, map2x, map2y, Q

def apply_rectification(left_img, right_img):
    """应用立体校正到左右图像"""
    ensure_initialized()
    # 1. 重映射
    rectified_left = cv2.remap(left_img, _map1x, _map1y, cv2.INTER_LINEAR)
    rectified_right = cv2.remap(right_img, _map2x, _map2y, cv2.INTER_LINEAR)
    
    return rectified_left, rectified_right

def calculate_fx_fy_cx_cy_baseline(Q):
    """计算fx,fy,cx,cy,baseline"""
    global _fx, _fy, _cx, _cy, _baseline
    
    Q = np.array(Q)
    _cx = -Q[0, 3]
    _cy = -Q[1, 3]
    _fx = Q[2, 3]
    _fy = Q[2, 3]
    _baseline = abs(1.0 / Q[3, 2]) * 1000

def reset_calib(scale_x, scale_y):
    """根据缩放比例重新计算内参"""
    global _fx, _fy, _cx, _cy, _depth_cam_matrix
    
    _fx, _fy = _fx * scale_x, _fy * scale_y
    _cx, _cy = _cx * scale_x, _cy * scale_y
    _depth_cam_matrix = np.array([[_fx, 0, _cx],
                                 [0, _fy, _cy],
                                 [0, 0, 1]])

def apply_rectification_single(left_img):
    """应用校正到单个图像"""
    ensure_initialized()
    # 重映射
    rectified_left = cv2.remap(left_img, _map1x, _map1y, cv2.INTER_LINEAR)
    return rectified_left

def load_calibration_from_json(file_path):
    """从JSON文件加载校准参数"""
    with open(file_path, 'r') as json_file:
        calibration_data = json.load(json_file)
    k_left = calibration_data['k_left']
    dist_left = calibration_data['dist_left']
    k_right = calibration_data['k_right']
    dist_right = calibration_data['dist_right']
    R = calibration_data['R']
    T = calibration_data['T']
    Q = calibration_data['Q']
    return k_left, dist_left, k_right, dist_right, R, T, Q

def rectified_image(left_image, right_image):
    """校正左右图像对应于原类的rectified_image方法"""
    return apply_rectification(left_image, right_image)

def disparity_to_depth(disp, focal_length, baseline):
    """视差图转换为深度图"""
    depth = (focal_length * baseline) / (disp + 1e-8)
    return depth

def ensure_initialized(resolution=(640, 480)):
    """确保已初始化，未初始化则进行初始化"""
    global _initialized
    if not _initialized:
        initialize(resolution)

# 导出的属性，与原类保持一致
def get_fx():
    ensure_initialized()
    return _fx

def get_fy():
    ensure_initialized()
    return _fy

def get_cx():
    ensure_initialized()
    return _cx

def get_cy():
    ensure_initialized()
    return _cy

def get_baseline():
    ensure_initialized()
    return _baseline

# 提供与原类相似的属性访问方式
class StereoRectifyCompat:
    """兼容层,提供与原stereoRectify类相似的接口"""
    def __init__(self, resolution=(640, 480)):
        initialize(resolution)
    
    @property
    def fx(self):
        return _fx
        
    @property
    def fy(self):
        return _fy
        
    @property
    def cx(self):
        return _cx
        
    @property
    def cy(self):
        return _cy
    def cali_cx(self):
        return _cx
        
    def cali_cy(self):
        return _cy
    @property
    def baseline(self):
        return _baseline
    
    @property
    def map1x(self):
        return _map1x
        
    @property
    def map1y(self):
        return _map1y
        
    @property
    def map2x(self):
        return _map2x
        
    @property
    def map2y(self):
        return _map2y
    
    def disparity_to_depth(self, disp, focal_length=None, baseline=None):
        if focal_length is None:
            focal_length = self.fx
        if baseline is None:
            baseline = self.baseline
        return disparity_to_depth(disp, focal_length, baseline)
        
    def apply_rectification(self, left_img, right_img, *args):
        return apply_rectification(left_img, right_img)
        
    def apply_rectification_single(self, left_img, *args):
        return apply_rectification_single(left_img)
        
    def rectified_image(self, left_image, right_image):
        return rectified_image(left_image, right_image)

# 为了向后兼容，初始化一个兼容对象
stereo_rectify = StereoRectifyCompat()

# 导出fx, fy, cx, cy, baseline作为模块属性，使其可以直接访问
fx = get_fx()
fy = get_fy()
cx = get_cx()
cy = get_cy()
baseline = get_baseline()

# 预初始化
initialize()


